import os
import struct
import shutil
from PIL import Image
from os import path

def classify_data(file,src,dir):
	counter=[0,0,0,0,0,0,0,0,0,0]
	with open(file,'r') as file:
		all=file.readline()
		labels=all.split(",")
		size=len(labels)
		for i in xrange(size):
			label=int(labels[i])
			counter[label] = counter[label]+1
			label_str=str(label)
			label_dir=dir+"/"+label_str
			if not path.exists(label_dir):
				os.mkdir(label_dir)
			src_file_name=src+"/"+str(i)+".png"
			dist_file_name=label_dir+"/"+str(i)+".png"
			shutil.copy(src_file_name,dist_file_name)
	print(counter)

def release_image(filename,dir):
	f=open(filename,'rb')
	index=0
	buf=f.read()
	f.close()

	magic,images,rows,columns=struct.unpack_from('>IIII',buf,index)
	index+=struct.calcsize('>IIII')
	for i in xrange(images):
		image = Image.new('L',(columns,rows))
		for x in xrange(rows):
			for y in xrange(columns):
				image.putpixel((y,x),int(struct.unpack_from('>B',buf,index)[0]))
				index += struct.calcsize('>B')
			image.save(dir+'/'+str(i)+'.png')

def init_dir():
	os.mkdir("train_data")
	os.mkdir("test_data")
	os.mkdir("label_data")
	os.mkdir("classified_train_data")
	os.mkdir("classified_test_data")

def release_label(filename,file):
	f=open(filename,'rb')
	index=0
	buf=f.read()
	f.close()
	magic,labels=struct.unpack_from('>II',buf,index)
	labelArr = [0] * labels
	index+=struct.calcsize('>II')
	for x in xrange(labels):
		labelArr[x]=int(struct.unpack_from('>B',buf,index)[0])
		index += struct.calcsize('>B')
	save=open(file,'w')
	save.write(','.join(map(lambda x: str(x), labelArr)))
	save.write('\n')
	save.close()
	print('save labels success')

if __name__ == "__main__":
	init_dir()
	release_image("original_data/train-images.idx3-ubyte","train_data")
	release_label("original_data/train-labels.idx1-ubyte","label_data/label_train.txt")
	release_image("original_data/t10k-images.idx3-ubyte","test_data")
	release_label("original_data/t10k-labels.idx1-ubyte","label_data/label_test.txt")
	classify_data("label_data/label_train.txt","train_data","classified_train_data")
	classify_data("label_data/label_test.txt","test_data","classified_test_data")